package 数组题目.前缀和与差分;

import java.util.HashMap;

//和可被 K 整除的子数组
//https://leetcode.cn/problems/subarray-sums-divisible-by-k/description/
class 和可被K整除的子数组 {
    //给定一个整数数组 nums 和一个整数 k ，返回其中元素之和可被 k 整除的非空 子数组 的数目。
    //子数组 是数组中 连续 的部分。
    public int subarraysDivByK(int[] nums, int k) {
        //(preSum[i] - preSum[j]) % k == 0
        // preSum[i] % k == preSum[j] % k
        // 0 4 9 9 7 4 5
        int[] preSum = new int[nums.length + 1];
        //存储 preSum[i] % k 的个数
        HashMap<Integer, Integer> indexCount = new HashMap<>();
        indexCount.put(0, 1);
        int res = 0;
        for (int i = 1; i < preSum.length; i++) {
            // 计算 nums[0..i] 的前缀和
            preSum[i] = preSum[i - 1] + nums[i - 1];
            // nums[0..i] 的所有元素之和与 k 的余数
            int index = preSum[i] % k;
            if (index < 0) {
                // 考虑到 preSum[i + 1] 可能是负数，根据 Java 求模的特性，-2 % 3 的结果是 -2
                // 但我们实际想要正余数 2，所以这里需要对 curRemainder 进行调整
                index += k;
            }
            // 看看之前 nums[0..i-1] 中是否也存在前缀和余数为 curRemainder 的子数组
            if (indexCount.containsKey(index)) {
                // 如果存在，则说明找到了可以整除 k 的子数组，累加子数组数量
                int count = indexCount.get(index);
                res += count;
                indexCount.put(index, count + 1);
            } else {
                // 如果不存在，那么 nums[0..i] 是第一个前缀和余数为 curRemainder 的子数组
                indexCount.put(index, 1);
            }
        }
        return res;
    }

    public static void main(String[] args) {
        System.out.println(new 和可被K整除的子数组().subarraysDivByK(new int[]{4, 5, 0, -2, -3, 1}, 5));
    }
}
